from langchain.prompts.prompt import PromptTemplate
from langchain.prompts.example_selector import LengthBasedExampleSelector
from langchain.prompts.few_shot import FewShotPromptTemplate
from langchain_openai import OpenAI

examples = [
    {"input": "高", "output": "矮"},
    {"input": "胖", "output": "瘦"},
    {"input": "精力充沛", "output": "萎靡不振"},
    {"input": "快乐", "output": "伤心"},
    {"input": "黑", "output": "白"},
]

example_prompt = PromptTemplate(input_variables=["input", "output"], template="""
词语：{input}\n
反义词：{output}\n
""")

example_selector = LengthBasedExampleSelector(
    examples=examples,
    example_prompt=example_prompt,
    max_length=25
)

few_shot_prompt = FewShotPromptTemplate(
    example_selector=example_selector,
    example_prompt=example_prompt,
    example_separator="\n",
    prefix="来玩个反义词接龙游戏，我说词语，你说它的反义词\n",
    suffix="现在轮到你了，词语：{input}\n反义词：",
    input_variables=["input"],
)

chat = OpenAI(temperature=0.3)
result = chat.invoke(few_shot_prompt.format(input="好"))
print(result)